package me.lwn.auth.security.oauth2.server;

import me.lwn.auth.security.provisioning.token.store.JdkSerializationStrategy;
import me.lwn.auth.security.utils.Constants;
import org.springframework.data.redis.connection.RedisConnection;
import org.springframework.data.redis.connection.RedisConnectionFactory;
import org.springframework.lang.Nullable;
import org.springframework.security.oauth2.core.OAuth2AccessToken;
import org.springframework.security.oauth2.core.OAuth2AuthorizationCode;
import org.springframework.security.oauth2.core.OAuth2RefreshToken;
import org.springframework.security.oauth2.core.OAuth2TokenType;
import org.springframework.security.oauth2.core.endpoint.OAuth2ParameterNames;
import org.springframework.security.oauth2.server.authorization.OAuth2Authorization;
import org.springframework.security.oauth2.server.authorization.OAuth2AuthorizationService;
import org.springframework.security.oauth2.server.authorization.client.RegisteredClient;
import org.springframework.security.oauth2.server.authorization.client.RegisteredClientRepository;
import org.springframework.util.ClassUtils;
import org.springframework.util.ReflectionUtils;
import org.springframework.util.StringUtils;

import java.lang.reflect.Method;
import java.time.Duration;
import java.time.Instant;
import java.util.Objects;

public class RedisOAuth2AuthorizationService implements OAuth2AuthorizationService {

    private static final String ACCESS = "access:";
    private static final String AUTH_TO_ACCESS = "auth_to_access:";
    private static final String AUTH = "auth:";
    private static final String REFRESH_AUTH = "refresh_auth:";
    private static final String ACCESS_TO_REFRESH = "access_to_refresh:";
    private static final String REFRESH = "refresh:";
    private static final String REFRESH_TO_ACCESS = "refresh_to_access:";
    private static final String CLIENT_ID_TO_ACCESS = "client_id_to_access:";
    private static final String UNAME_TO_ACCESS = "uname_to_access:";
    private static final String CODE_TO_AUTH = "code_to_auth:";
    private static final String AUTH_TO_CODE = "auth_to_code:";
    private static final String STATE_TO_AUTH = "state_to_auth:";
    private static final String AUTH_TO_STATE = "auth_to_state:";
    private static final boolean springDataRedis_2_0 = ClassUtils.isPresent(
            "org.springframework.data.redis.connection.RedisStandaloneConfiguration",
            RedisOAuth2AuthorizationService.class.getClassLoader());
    private final RedisConnectionFactory connectionFactory;
    private final JdkSerializationStrategy serializationStrategy = new JdkSerializationStrategy();
    private final RegisteredClientRepository repository;
    private Method redisConnectionSet_2_0;
    private String prefix = "";

    public RedisOAuth2AuthorizationService(RedisConnectionFactory connectionFactory, RegisteredClientRepository repository) {
        this.connectionFactory = connectionFactory;
        this.repository = repository;
        if (springDataRedis_2_0) {
            this.loadRedisConnectionMethods_2_0();
        }
    }

    private static String getApprovalKey(OAuth2Authorization authorization) {
        String userName = authorization.getPrincipalName() == null ? ""
                : authorization.getPrincipalName();
        return getApprovalKey(authorization.getRegisteredClientId(), userName);
    }

    private static String getApprovalKey(String clientId, String userName) {
        return clientId + (userName == null ? "" : ":" + userName);
    }

    private void loadRedisConnectionMethods_2_0() {
        this.redisConnectionSet_2_0 = ReflectionUtils.findMethod(
                RedisConnection.class, "set", byte[].class, byte[].class);
    }

    private RedisConnection getConnection() {
        return connectionFactory.getConnection();
    }

    private byte[] serialize(Object obj) {
        return serializationStrategy.serialize(obj);
    }

    private byte[] serialize(String str) {
        return serializationStrategy.serialize(str);
    }

    private byte[] serializeKey(String key) {
        return serialize(prefix + key);
    }

    private OAuth2AccessToken deserializeOAuth2AccessToken(byte[] bytes) {
        return serializationStrategy.deserialize(bytes, OAuth2AccessToken.class);
    }

    private OAuth2Authorization deserializeOAuth2Authorization(byte[] bytes) {
        return serializationStrategy.deserialize(bytes, OAuth2Authorization.class);
    }

    private void storeAccessToken(RedisConnection connection, OAuth2Authorization authorization) {

        OAuth2AccessToken token = authorization.getAccessToken().getToken();
        byte[] serializedAccessToken = serialize(token);
        byte[] serializedAuth = serialize(authorization);
        byte[] accessKey = serializeKey(ACCESS + token.getTokenValue());
        byte[] authKey = serializeKey(AUTH + token.getTokenValue());
        byte[] authToAccessKey = serializeKey(AUTH_TO_ACCESS + authorization.getId());
        byte[] approvalKey = serializeKey(UNAME_TO_ACCESS + getApprovalKey(authorization));
        byte[] clientId = serializeKey(CLIENT_ID_TO_ACCESS + authorization.getRegisteredClientId());
        if (springDataRedis_2_0) {
            try {
                this.redisConnectionSet_2_0.invoke(connection, accessKey, serializedAccessToken);
                this.redisConnectionSet_2_0.invoke(connection, authKey, serializedAuth);
                this.redisConnectionSet_2_0.invoke(connection, authToAccessKey, serializedAccessToken);
            } catch (Exception ex) {
                throw new RuntimeException(ex);
            }
        } else {
            connection.set(accessKey, serializedAccessToken);
            connection.set(authKey, serializedAuth);
            connection.set(authToAccessKey, serializedAccessToken);
        }
        if (StringUtils.hasText(authorization.getPrincipalName())) {
            connection.sAdd(approvalKey, serializedAccessToken);
        }
        connection.sAdd(clientId, serializedAccessToken);
        if (token.getExpiresAt() != null) {
            long seconds = Duration.between(Instant.now(), token.getExpiresAt()).getSeconds();
            connection.expire(accessKey, seconds);
            connection.expire(authKey, seconds);
            connection.expire(authToAccessKey, seconds);
            connection.expire(clientId, seconds);
            connection.expire(approvalKey, seconds);
        }
    }

    private void storeRefreshToken(RedisConnection connection, OAuth2Authorization authorization) {
        OAuth2AccessToken token = authorization.getAccessToken().getToken();
        if (authorization.getRefreshToken() != null) {
            OAuth2RefreshToken refreshToken = authorization.getRefreshToken().getToken();
            if (refreshToken != null && refreshToken.getTokenValue() != null) {
                byte[] refresh = serialize(refreshToken.getTokenValue());
                byte[] auth = serialize(token.getTokenValue());
                byte[] refreshKey = serializeKey(REFRESH + refreshToken.getTokenValue());
                byte[] refreshAuthKey = serializeKey(REFRESH_AUTH + refreshToken.getTokenValue());
                byte[] serializedRefreshToken = serialize(refreshToken);
                byte[] refreshToAccessKey = serializeKey(REFRESH_TO_ACCESS + refreshToken.getTokenValue());
                byte[] accessToRefreshKey = serializeKey(ACCESS_TO_REFRESH + token.getTokenValue());
                if (springDataRedis_2_0) {
                    try {
                        this.redisConnectionSet_2_0.invoke(connection, refreshToAccessKey, auth);
                        this.redisConnectionSet_2_0.invoke(connection, accessToRefreshKey, refresh);
                        this.redisConnectionSet_2_0.invoke(connection, refreshKey, serializedRefreshToken);
                        this.redisConnectionSet_2_0.invoke(connection, refreshAuthKey, serialize(authorization));
                    } catch (Exception ex) {
                        throw new RuntimeException(ex);
                    }
                } else {
                    connection.set(refreshToAccessKey, auth);
                    connection.set(accessToRefreshKey, refresh);
                    connection.set(refreshKey, serializedRefreshToken);
                    connection.set(refreshAuthKey, serialize(authorization));
                }
                if (refreshToken.getExpiresAt() != null) {
                    long seconds = Duration.between(Instant.now(), refreshToken.getExpiresAt()).getSeconds();
                    connection.expire(refreshToAccessKey, seconds);
                    connection.expire(accessToRefreshKey, seconds);
                    connection.expire(refreshKey, seconds);
                    connection.expire(refreshAuthKey, seconds);
                }
            }
        }
    }

    private void storeOAuthCode(RedisConnection connection, OAuth2Authorization authorization, OAuth2AuthorizationCode authorizationCode) {
        byte[] codeToAuthKey = serializeKey(CODE_TO_AUTH + authorizationCode.getTokenValue());
        byte[] authToCodeKey = serializeKey(AUTH_TO_CODE + authorization.getId());

        if (springDataRedis_2_0) {
            try {
                this.redisConnectionSet_2_0.invoke(connection, codeToAuthKey, serialize(authorization));
                this.redisConnectionSet_2_0.invoke(connection, authToCodeKey, serialize(authorization.getId()));
            } catch (Exception ex) {
                throw new RuntimeException(ex);
            }
        } else {
            connection.set(codeToAuthKey, serialize(authorization));
            connection.set(authToCodeKey, serialize(authorization.getId()));
        }

        if (authorizationCode.getExpiresAt() != null) {
            long seconds = Duration.between(Instant.now(), authorizationCode.getExpiresAt()).getSeconds();
            connection.expire(codeToAuthKey, seconds);
            connection.expire(authToCodeKey, seconds);
        }
    }

    private void storeState(RedisConnection connection, OAuth2Authorization authorization) {

        String state = authorization.getAttribute(OAuth2ParameterNames.STATE);
        byte[] stateToAuthKey = serializeKey(STATE_TO_AUTH + state);
        byte[] authToStateKey = serializeKey(AUTH_TO_STATE + authorization.getId());

        if (springDataRedis_2_0) {
            try {
                this.redisConnectionSet_2_0.invoke(connection, stateToAuthKey, serialize(authorization));
                this.redisConnectionSet_2_0.invoke(connection, authToStateKey, serialize(state));
            } catch (Exception ex) {
                throw new RuntimeException(ex);
            }
        } else {
            connection.set(stateToAuthKey, serialize(authorization));
            connection.set(authToStateKey, serialize(state));
        }

        RegisteredClient client = repository.findByClientId(authorization.getRegisteredClientId());
        Long duration = (Long) Objects.requireNonNull(client).getTokenSettings().getSettings().get(Constants.TokenSetting.STATE_TIME_TO_LIVE.getKey());
        if (duration == null) {
            duration = 300L;
        }
        connection.expire(stateToAuthKey, duration);
        connection.expire(authToStateKey, duration);
    }

    private boolean stateIsExist(RedisConnection connection, String id) {
        return findBy(connection, AUTH_TO_STATE + id, String.class) != null;
    }

    @Override
    public void save(OAuth2Authorization authorization) {

        try (RedisConnection conn = getConnection()) {
            conn.openPipeline();
            String state = authorization.getAttribute(OAuth2ParameterNames.STATE);
            OAuth2Authorization.Token<OAuth2AuthorizationCode> code = authorization.getToken(OAuth2AuthorizationCode.class);

            if (code != null && code.getToken() != null) {
                if (code.isActive()) {
                    storeOAuthCode(conn, authorization, code.getToken());
                } else {
                    removeOAuthCode(conn, authorization.getId(), code.getToken());
                }
            }

            if (stateIsExist(conn, authorization.getId())) {
                removeState(conn, authorization.getId());
            }

            if (StringUtils.hasText(state)) {
                storeState(conn, authorization);
            }

            if (authorization.getToken(OAuth2AccessToken.class) != null || authorization.getToken(OAuth2RefreshToken.class) != null) {

                OAuth2Authorization auth = findById(authorization.getId());
                if (auth != null) {
                    remove(conn, auth);
                }

                storeAccessToken(conn, authorization);
                storeRefreshToken(conn, authorization);
            }
            conn.closePipeline();
        }
    }

    public void removeOAuthCode(RedisConnection conn, String id, OAuth2AuthorizationCode authCode) {
        byte[] codeToAuthKey = serializeKey(CODE_TO_AUTH + authCode.getTokenValue());
        byte[] authToCodeKey = serializeKey(AUTH_TO_CODE + id);
        conn.del(codeToAuthKey);
        conn.del(authToCodeKey);
    }

    public void removeState(RedisConnection conn, String id) {
        byte[] authToStateKey = serializeKey(AUTH_TO_STATE + id);
        byte[] states = conn.get(serializeKey(AUTH_TO_ACCESS + id));
        String state = serializationStrategy.deserialize(states, String.class);
        conn.del(serializeKey(STATE_TO_AUTH + state));
        conn.del(authToStateKey);
    }

    public void remove(RedisConnection conn, OAuth2Authorization authorization) {
        OAuth2AccessToken token = authorization.getAccessToken().getToken();
        byte[] serializedAccessToken = serialize(token);
        byte[] accessKey = serializeKey(ACCESS + token.getTokenValue());
        byte[] authKey = serializeKey(AUTH + token.getTokenValue());
        byte[] authToAccessKey = serializeKey(AUTH_TO_ACCESS + authorization.getId());
        byte[] approvalKey = serializeKey(UNAME_TO_ACCESS + getApprovalKey(authorization));
        byte[] clientId = serializeKey(CLIENT_ID_TO_ACCESS + authorization.getRegisteredClientId());

        conn.del(accessKey);
        conn.del(authKey);
        conn.del(authToAccessKey);
        conn.sRem(approvalKey, serializedAccessToken);
        conn.sRem(clientId, serializedAccessToken);

        if (authorization.getRefreshToken() != null) {
            OAuth2RefreshToken refreshToken = authorization.getRefreshToken().getToken();
            if (refreshToken != null && refreshToken.getTokenValue() != null) {
                byte[] refreshKey = serializeKey(REFRESH + refreshToken.getTokenValue());
                byte[] refreshAuthKey = serializeKey(REFRESH_AUTH + refreshToken.getTokenValue());
                byte[] refreshToAccessKey = serializeKey(REFRESH_TO_ACCESS + refreshToken.getTokenValue());
                byte[] accessToRefreshKey = serializeKey(ACCESS_TO_REFRESH + token.getTokenValue());
                conn.del(refreshKey);
                conn.del(refreshAuthKey);
                conn.del(refreshToAccessKey);
                conn.del(accessToRefreshKey);
            }
        }
    }

    @Override
    public void remove(OAuth2Authorization authorization) {
        // TODO 此处应该根据 OAuth2Authorization 执行不同的删除逻辑。暂时不强制使用此删除方法
        try (RedisConnection conn = getConnection()) {
            if (StringUtils.hasText(authorization.getAttribute(OAuth2ParameterNames.STATE))) {
                removeState(conn, authorization.getAttribute(OAuth2ParameterNames.STATE));
            }
            remove(conn, authorization);
        }
        //access               OAuth2AccessToken
        //access_to_refresh    OAuth2RefreshToken
        //access_to_grant      OAuth2Authentication
        //auth                 OAuth2Authentication
        //auth_to_access       OAuth2AccessToken
        //refresh              OAuth2RefreshToken
        //refresh_to_auth      OAuth2Authentication
        //refresh_to_access    access_token
    }

    @Nullable
    @Override
    public OAuth2Authorization findById(String id) {
        byte[] bytes;
        try (RedisConnection conn = getConnection()) {
            bytes = conn.get(serializeKey(AUTH_TO_ACCESS + id));
            OAuth2AccessToken accessToken = deserializeOAuth2AccessToken(bytes);
            if (accessToken == null) {
                return null;
            }
            bytes = conn.get(serializeKey(AUTH + accessToken.getTokenValue()));
        }
        return deserializeOAuth2Authorization(bytes);
    }

    @Nullable
    @Override
    public OAuth2Authorization findByToken(String token, @Nullable OAuth2TokenType tokenType) {
        if (tokenType == null) {
            return null;
        } else if (OAuth2ParameterNames.STATE.equals(tokenType.getValue())) {
            return findBy(STATE_TO_AUTH + token);
        } else if (OAuth2ParameterNames.CODE.equals(tokenType.getValue())) {
            return findBy(CODE_TO_AUTH + token);
        } else if (OAuth2TokenType.ACCESS_TOKEN.equals(tokenType)) {
            return findBy(AUTH + token);
        } else if (OAuth2TokenType.REFRESH_TOKEN.equals(tokenType)) {
            return findBy(REFRESH_AUTH + token);
        }
        return null;
    }

    public <T> T findBy(RedisConnection connection, String key, Class<T> t) {
        byte[] bytes = connection.get(serializeKey(key));
        return serializationStrategy.deserialize(bytes, t);
    }

    public OAuth2Authorization findBy(String key) {
        byte[] bytes;
        try (RedisConnection conn = getConnection()) {
            bytes = conn.get(serializeKey(key));
        }
        return deserializeOAuth2Authorization(bytes);
    }

    public void setPrefix(String prefix) {
        this.prefix = prefix;
    }
}
